# -*- coding: UTF-8 -*-
#
import numpy as np
import scipy.fftpack as fft
import math
#
def calcElasticWavePhysicalParams(rho, E, nu):
	G = E/(2*(1+nu))			# shear modulus
	lamb = 2*nu*G/(1-2*nu)		# lame constant
	#
	vs = math.sqrt(G/rho)			# s-wave velocity
	kps = math.sqrt((2-2*nu)/(1-2*nu))			# kps = vp/vs
	krs = -0.07602*(nu**2)+0.2012*nu+0.8738	# krs = vr/vs
	krp = krs/kps		# krp = vr/vp
	vp = vs*kps			# p-wave velocity
	vr = vs*krs			# Rayleigh wave velocity
	#
	return (G,lamb,vs,vp,vr,krp,krs)

def genImpulseWave(timePointData, Tp, Amp):
	assert (Amp>0) and (timePointData[-1]>=Tp)
	heaviside = lambda t : 0.5*(np.sign(t)+1.0)
	G4 = lambda t : np.power(t,3)*heaviside(t)
	delta = lambda t : 16*(G4(t)-4*G4(t-0.25)+6*G4(t-0.5)-4*G4(t-0.75)+G4(t-1.0))
	impulse = lambda t,Td,A : A*delta(t/Td)*(t<=Td)
	return impulse(timePointData, Tp, Amp)

def getArrayAbsMaxPos(array):
	absMaxPos = 0
	for i in range(array.shape[0]):
		if math.fabs(array[absMaxPos])<math.fabs(array[i]):
			absMaxPos = i
	return absMaxPos

def genRayleighInputWave(timePointData, Tp, freq, Amp):
	envelope = genImpulseWave(timePointData, Tp, Amp)
	omega = 2.0*np.pi*freq
	ux = envelope*np.cos(omega*timePointData)
	ux = Amp*ux/ux[getArrayAbsMaxPos(ux)]
	return ux

def calcRayleighWaveParams(ux, dt, vr, krp, krs):
	N = ux.shape[0]
	#assert (N%2)==0
	#
	a = math.sqrt(1-krp**2)
	b = math.sqrt(1-krs**2)
	domega = np.pi/((N+1)*dt)
	omega = np.arange(1.0,N+1.0)*domega
	k = omega/vr
	#
	coef = fft.dst(ux,1)/(N+1)
	z = 0
	eakz = np.exp(-a*k*z)
	ebkz = np.exp(-b*k*z)
	f1 = k*(eakz-2*a*b/(1+b**2)*ebkz)
	A = coef/f1
	return (a,b,k,omega,A)

def calcRayleighWaveF1F2(a, b, k, z):
	eakz = np.exp(-a*k*z)
	ebkz = np.exp(-b*k*z)
	f1 = k*( eakz-2*a*b/(1+b**2)*ebkz)
	f2 = k*(-a*eakz+2*a/(1+b**2)*ebkz)
	return (f1, f2)

def calcRayleighWaveF1F2F3F4F5(a, b, k, z):
	eakz = np.exp(-a*k*z)
	ebkz = np.exp(-b*k*z)
	f1 = k*( eakz-2*a*b/(1+b**2)*ebkz)
	f2 = k*(-a*eakz+2*a/(1+b**2)*ebkz)
	f3 = (k**2)*(-eakz+2*a*b/(1+b**2)*ebkz)
	f4 = (k**2)*((a**2)*eakz-2*a*b/(1+b**2)*ebkz)
	f5 = (k**2)*(-2*a*eakz+2*a*ebkz)
	return (f1, f2, f3, f4, f5)

def calcRayleighWaveAccDisp(A,f1,f2,omega,t,x,vr):
	N = A.shape[0]
	uxmm = np.zeros(shape=(N,N))
	uzmm = np.zeros(shape=(N,N))
	axmm = np.zeros(shape=(N,N))
	azmm = np.zeros(shape=(N,N))
	#
	txvr = t-x/vr
	txvr_mask = (txvr>=0)
	for i in range(N):
		swit = np.sin(omega[i]*txvr)*txvr_mask		# sin(omega*t)
		cwit = np.cos(omega[i]*txvr)*txvr_mask		# cos(omega*t)
		uxmm[i,:] =  A[i]*f1[i]*swit
		uzmm[i,:] =  A[i]*f2[i]*cwit
		axmm[i,:] = -A[i]*f1[i]*omega[i]*omega[i]*swit
		azmm[i,:] = -A[i]*f2[i]*omega[i]*omega[i]*cwit
	ux = np.sum(uxmm, axis=0)
	uz = np.sum(uzmm, axis=0)
	ax = np.sum(axmm, axis=0)
	az = np.sum(azmm, axis=0)
	return (ax,az,ux,uz)

def calcRayleighWaveDispVelStrain(A,f1,f2,f3,f4,f5,omega,t,x,vr):
	N = A.shape[0]
	uxmm  = np.zeros(shape=(N,N))
	uzmm  = np.zeros(shape=(N,N))
	vxmm  = np.zeros(shape=(N,N))
	vzmm  = np.zeros(shape=(N,N))
	exxmm = np.zeros(shape=(N,N))
	ezzmm = np.zeros(shape=(N,N))
	exzmm = np.zeros(shape=(N,N))
	#
	txvr = t-x/vr
	txvr_mask = (txvr>=0)
	for i in range(N):
		swit = np.sin(omega[i]*txvr)*txvr_mask		# sin(omega*t)
		cwit = np.cos(omega[i]*txvr)*txvr_mask		# cos(omega*t)
		uxmm[i,:]  =  A[i]*f1[i]*swit
		uzmm[i,:]  =  A[i]*f2[i]*cwit
		vxmm[i,:]  =  A[i]*f1[i]*omega[i]*cwit
		vzmm[i,:]  = -A[i]*f2[i]*omega[i]*swit
		exxmm[i,:] =  A[i]*f3[i]*cwit
		ezzmm[i,:] =  A[i]*f4[i]*cwit
		exzmm[i,:] =  A[i]*f5[i]*swit
	ux = np.sum(uxmm, axis=0)
	uz = np.sum(uzmm, axis=0)
	vx = np.sum(vxmm, axis=0)
	vz = np.sum(vzmm, axis=0)
	exx = np.sum(exxmm, axis=0)
	ezz = np.sum(ezzmm, axis=0)
	exz = np.sum(exzmm, axis=0)
	return (ux,uz,vx,vz,exx,ezz,exz)

